package org.zoomdev.zoom.lockfree;

import java.lang.ScopedValue;
import java.lang.invoke.MethodHandles;
import java.lang.invoke.VarHandle;
import java.util.concurrent.ExecutorService;
import java.util.concurrent.Executors;
import java.util.function.Supplier;

public class LocklessObjectPool<T> {
    // 栈节点实现
    private static class Node<T> {
        volatile Node<T> next;
        final T object;
        final long createTime;

        Node(T object) {
            this.object = object;
            this.createTime = System.currentTimeMillis();
        }
    }

    // 虚拟线程本地存储
    private static final ScopedValue<Node<?>> CURRENT_NODE = ScopedValue.newInstance();

    // 无锁栈顶指针
    private volatile Node<T> stack;

    // VarHandle实现无锁操作
    private static final VarHandle STACK_HANDLE;
    static {
        try {
            MethodHandles.Lookup lookup = MethodHandles.lookup();
            STACK_HANDLE = lookup.findVarHandle(
                    LocklessObjectPool.class, "stack", Node.class);
        } catch (Exception e) {
            throw new Error(e);
        }
    }

    private final Supplier<T> objectFactory;
    private final int minSize;
    private final int maxSize;
    private volatile int currentSize;

    public LocklessObjectPool(Supplier<T> objectFactory,
                              int minSize,
                              int maxSize) {
        this.objectFactory = objectFactory;
        this.minSize = minSize;
        this.maxSize = maxSize;
        initializePool();
    }

    private void initializePool() {
        for (int i = 0; i < minSize; i++) {
            T obj = objectFactory.get();
            Node<T> node = new Node<>(obj);
            push(node);
            currentSize++;
        }
    }

    // 压栈操作 (CAS实现)
    private void push(Node<T> node) {
        Node<T> oldHead;
        do {
            oldHead = stack;
            node.next = oldHead;
        } while (!STACK_HANDLE.compareAndSet(this, oldHead, node));
    }

    // 出栈操作 (CAS实现)
    private Node<T> pop() {
        Node<T> oldHead;
        Node<T> newHead;

        do {
            oldHead = stack;
            if (oldHead == null) {
                return null;
            }
            newHead = oldHead.next;
        } while (!STACK_HANDLE.compareAndSet(this, oldHead, newHead));

        return oldHead;
    }

    // 借出对象
    public T borrow() {
        // 检查当前虚拟线程是否已持有对象
        if (CURRENT_NODE.isBound()) {
            throw new IllegalStateException("Recursive borrow not allowed");
        }

        Node<T> node = pop();

        if (node == null) {
            // 池中无可用对象，考虑创建新对象
            if (currentSize < maxSize) {
                // CAS增加当前池大小
                int current;
                do {
                    current = currentSize;
                    if (current >= maxSize) {
                        // 其他线程已扩容
                        node = pop();
                        break;
                    }
                } while (!compareAndSetCurrentSize(current, current + 1));

                if (node == null) {
                    T newObj = objectFactory.get();
                    node = new Node<>(newObj);
                }
            } else {
                // 池已达最大值，等待释放
                return waitForAvailable();
            }
        }

        final Node<T> acquiredNode = node;
        // 绑定对象到虚拟线程范围
        ScopedValue.where(CURRENT_NODE, acquiredNode).run(() -> {
            // 可在此处执行对象初始化逻辑
        });

        return acquiredNode.object;
    }

    // 归还对象
    public void release() {
        Node<?> node = CURRENT_NODE.get();
        if (node == null) {
            throw new IllegalStateException("No object acquired in current scope");
        }

        if (currentSize > minSize) {
            // 检查对象存活时间
            long liveTime = System.currentTimeMillis() - node.createTime;
            if (liveTime > 300_000) { // 5分钟
                // 对象过旧，销毁
                destroyObject((Node<T>) node);
                return;
            }
        }

        // 对象入栈复用
        push((Node<T>) node);
    }

    // 等待可用对象（无锁实现）
    private T waitForAvailable() {
        // 虚拟线程友好等待算法
        return VirtualThreadExecutor.execute(() -> {
            long start = System.nanoTime();
            int attempt = 0;

            while (true) {
                Node<T> node = pop();
                if (node != null) {
                    return node.object;
                }

                // 指数退避算法
                long sleepNanos = Math.min(100_000 * (1L << attempt), 10_000_000);
                try {
                    Thread.sleep(sleepNanos / 1_000_000, (int) (sleepNanos % 1_000_000));
                } catch (InterruptedException e) {
                    Thread.currentThread().interrupt();
                    throw new PoolException("Wait interrupted");
                }

                if (System.nanoTime() - start > 30_000_000_000L) { // 30秒超时
                    throw new PoolException("Timeout waiting for object");
                }

                attempt = Math.min(attempt + 1, 16);
            }
        });
    }

    // 销毁对象（无锁缩减）
    private void destroyObject(Node<T> node) {
        // 对象销毁前的清理逻辑
        if (node.object instanceof AutoCloseable) {
            try {
                ((AutoCloseable) node.object).close();
            } catch (Exception e) {
                // 日志记录
            }
        }

        // CAS减少池大小
        int current;
        do {
            current = currentSize;
            if (current <= minSize) break;
        } while (!compareAndSetCurrentSize(current, current - 1));
    }

    // VarHandle辅助方法
    private static final VarHandle SIZE_HANDLE;
    static {
        try {
            MethodHandles.Lookup lookup = MethodHandles.lookup();
            SIZE_HANDLE = lookup.findVarHandle(
                    LocklessObjectPool.class, "currentSize", int.class);
        } catch (Exception e) {
            throw new Error(e);
        }
    }

    private boolean compareAndSetCurrentSize(int expected, int update) {
        return SIZE_HANDLE.compareAndSet(this, expected, update);
    }

    // 虚拟线程专用执行器
    private static class VirtualThreadExecutor {
        private static final ExecutorService EXECUTOR =
                Executors.newThreadPerTaskExecutor(
                        Thread.ofVirtual().name("pool-waiter-", 0).factory()
                );

        static <U> U execute(Supplier<U> task) {
            try {
                return EXECUTOR.submit(task::get).get();
            } catch (Exception e) {
                throw new PoolException("Execution failed", e);
            }
        }
    }

    // 异常类
    public static class PoolException extends RuntimeException {
        public PoolException(String message) {
            super(message);
        }
        public PoolException(String message, Throwable cause) {
            super(message, cause);
        }
    }
}